from flask import Flask, render_template, request, jsonify
from flask_socketio import SocketIO, emit
import numpy as np
import json
import time
import threading

app = Flask(__name__)
app.config['SECRET_KEY'] = 'secret!'
socketio = SocketIO(app, cors_allowed_origins="*")

class FloodSimulation:
    def __init__(self):
        self.L = 1000  # 巷道长度 (m)
        self.B = 5     # 矩形巷道宽度 (m)
        self.n = 0.04  # 曼宁粗糙系数
        self.S0 = 0.002  # 巷道底坡
        self.g = 9.81  # 重力加速度
        
        # 突水参数
        self.Q_inject = 15  # 突水峰值流量 (m³/s)
        self.t_peak = 60    # 达到峰值流量时间 (s)
        self.t_total = 200  # 突水总持续时间 (s)
        self.x_inject = 300  # 突水发生位置 (m)
        
        # 数值计算参数
        self.Nx = 101
        self.Nt = 1000
        self.T_max = 600
        self.theta = 0.65
        
        self.is_running = False
        self.current_time = 0
        
    def calculate_flow_rate(self, t):
        """计算t时刻的突水流量"""
        if t <= 0:
            return 0
        elif t <= self.t_peak:
            return self.Q_inject * (t / self.t_peak)
        elif t <= self.t_total:
            return self.Q_inject * (1 - (t - self.t_peak) / (self.t_total - self.t_peak))
        else:
            return 0
    
    def saint_venant_solver(self, h0, Q0, dt, dx, t):
        """简化的圣维南方程求解器"""
        h = h0.copy()
        Q = Q0.copy()
        
        # 计算流量
        Q_in = self.calculate_flow_rate(t)
        
        # 简化的水流传播模型
        alpha = 0.1  # 扩散系数
        for i in range(1, self.Nx-1):
            # 水深扩散
            h[i] = h0[i] + alpha * dt / dx**2 * (h0[i+1] - 2*h0[i] + h0[i-1])
            
            # 在突水点添加流量
            x_pos = i * dx
            if abs(x_pos - self.x_inject) < dx:
                h[i] += Q_in * dt / (dx * self.B) * 0.1
        
        # 边界条件
        h[0] = h[1]
        h[-1] = h[-2]
        
        # 计算流速
        v = np.zeros_like(h)
        for i in range(self.Nx):
            if h[i] > 0.01:
                # 曼宁公式计算流速
                Rh = (self.B * h[i]) / (self.B + 2 * h[i])  # 水力半径
                v[i] = (1/self.n) * Rh**(2/3) * self.S0**0.5
        
        return h, v
    
    def run_simulation(self):
        """运行模拟"""
        self.is_running = True
        
        # 初始化
        dx = self.L / (self.Nx - 1)
        dt = self.T_max / self.Nt
        
        h = np.zeros(self.Nx)
        v = np.zeros(self.Nx)
        x = np.linspace(0, self.L, self.Nx)
        
        for step in range(self.Nt):
            if not self.is_running:
                break
                
            t = step * dt
            h, v = self.saint_venant_solver(h, h, dt, dx, t)
            
            # 发送数据到前端
            data = {
                'time': t,
                'x': x.tolist(),
                'h': h.tolist(),
                'v': v.tolist(),
                'Q': self.calculate_flow_rate(t)
            }
            
            socketio.emit('simulation_data', data)
            time.sleep(0.1)  # 控制发送频率
            
            self.current_time = t
        
        self.is_running = False
        socketio.emit('simulation_complete')

simulation = FloodSimulation()

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/api/start', methods=['POST'])
def start_simulation():
    if simulation.is_running:
        return jsonify({'error': 'Simulation already running'}), 400
    
    # 更新参数
    data = request.json
    simulation.L = float(data.get('L', 1000))
    simulation.B = float(data.get('B', 5))
    simulation.n = float(data.get('n', 0.04))
    simulation.S0 = float(data.get('S0', 0.002))
    simulation.Q_inject = float(data.get('Q_inject', 15))
    simulation.t_peak = float(data.get('t_peak', 60))
    simulation.t_total = float(data.get('t_total', 200))
    simulation.x_inject = float(data.get('x_inject', 300))
    
    # 启动模拟线程
    thread = threading.Thread(target=simulation.run_simulation)
    thread.start()
    
    return jsonify({'status': 'started'})

@app.route('/api/stop', methods=['POST'])
def stop_simulation():
    simulation.is_running = False
    return jsonify({'status': 'stopped'})

@app.route('/api/reset', methods=['POST'])
def reset_simulation():
    simulation.is_running = False
    simulation.current_time = 0
    return jsonify({'status': 'reset'})

if __name__ == '__main__':
    socketio.run(app, debug=True, host='0.0.0.0', port=5000)
